from osgeo import gdal, ogr
import os


# shapefile 转 geojson
def to_geojson(shapefile, output_json):
    # 打开矢量图层
    gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
    gdal.SetConfigOption("SHAPE_ENCODING", "GBK")
    shp_dataSource = ogr.Open(shapefile)
    shp_layer = shp_dataSource.GetLayer(0)

    # 创建Geojson图层
    base_name = os.path.basename(output_json)
    json_driver = ogr.GetDriverByName('GeoJSON')
    json_ds = json_driver.CreateDataSource(output_json)
    if json_ds.GetLayer(base_name):
        json_ds.DeleteLayer(base_name)  # 如果已有文件则删除
    # 创建json图层,json图层名，设置空间参考【同shp数据的空间参考一致】
    json_layer = json_ds.CreateLayer(base_name, shp_layer.GetSpatialRef())
    json_layer.CreateFields(shp_layer.schema)  # 创建属性字段
    # 创建要素
    json_feature = ogr.Feature(json_layer.GetLayerDefn())

    # 遍历shapefile图层中的所有要素，并将要素装载到json的图层的要素中
    for feature in shp_layer:
        json_feature.SetGeometry(feature.geometry())
        for j in range(feature.GetFieldCount()):
            # 遍历每一个要素的中的所有字段属性，并将字段属性装载到json的要素中
            json_feature.SetField(j, feature.GetField(j))
        # 把json要素装载到json图层中
        json_layer.CreateFeature(json_feature)

    del json_ds  # 释放内存
    del shp_dataSource  # 释放内存
    print("to geojson success")


if __name__ == '__main__':
    shapefile = 'data/Search_Polygon.shp'
    output_json = 'data/Search_Polygon.json'
    to_geojson(shapefile, output_json)
